package pl;

import javacard.framework.JCSystem;
import javacard.framework.Util;
import javacard.security.MessageDigest;

public class KDF2BytesGenerator {

	private short counterStart;
	private MessageDigest md;
	private byte[] shared;
	private byte[] iv;
	byte[] C;
	
	protected KDF2BytesGenerator(MessageDigest digest) {
		// default value for KDF2
		this.counterStart = 1;
		this.md = digest;
		C = JCSystem.makeTransientByteArray((short)4, JCSystem.CLEAR_ON_DESELECT); //new byte[4];
	}

	// derivatiom parameters -> shared
	public void init(byte[] shared) {
		this.shared = shared;
	}
	
	private void I2OSP(short i, byte[] sp) {
		sp[0] = (byte) (i >>> 24);
		sp[1] = (byte) (i >>> 16);
		sp[2] = (byte) (i >>> 8);
		sp[3] = (byte) (i >>> 0);
	}

	public short generateBytes(byte[] out, short outOff, short len)
	{
		if ((out.length - len) < outOff) {
			// throw new DataLengthException("output buffer too small");
		}

		short oBytes = len;
		// short outLen = md.getDigestSize();
		short outLen = md.getLength();
		//
		// this is at odds with the standard implementation, the
		// maximum value should be hBits * (2^32 - 1) where hBits
		// is the digest output size in bits. We can't have an
		// array with a long index at the moment...
		//
		if (oBytes > ((2L << 32) - 1)) {
			// throw new IllegalArgumentException("Output length too large");
		}

		short cThreshold = (short) (((short) (oBytes + outLen - 1)) / outLen);

		byte[] dig = null;

		// dig = new byte[digest.getDigestSize()];
		dig = new byte[md.getLength()];

		short counter = 0; // counterStart;

		for (short i = 0; i < cThreshold; i++) {
			// shared -> bos
			// md.update(shared, 0, shared.length);
			md.update(shared, (short) 0, (short) shared.length);

			I2OSP(counter, C);
			
			md.update(C, (short) 0, (short)4);
			
			md.doFinal(C, (short) 0, (short) C.length, dig, (short) 0);
			
			if (len > outLen) {
				// System.arraycopy(dig, 0, out, outOff, outLen);

				Util.arrayCopy(dig, (short) 0, out, outOff, outLen);
				outOff += outLen;
				len -= outLen;
			} else {
				Util.arrayCopy(dig, (short) 0, out, (short) 0, len);
				// System.arraycopy(dig, 0, out, outOff, len);
			}

			counter++;
		}

		// digest.reset();
		md.reset();
		
		return len;
	}
}
